#!/usr/bin/python2
"""
Read ROS bag message, then save it to folder
"""

import os
import argparse
import rosbag
import numpy as np
import cv2 as cv
import cv_bridge
import logging


def main(bag_file, save_folder):
    # open file and print topics
    logging.info('bag file: {}'.format(bag_file))
    logging.info('save folder: {}'.format(save_folder))

    # create save folder
    save_root = os.path.abspath(save_folder)
    if not os.path.exists(save_root):
        os.makedirs(save_root)
    # create image save folder
    img_folder = os.path.join(save_root, 'camera')
    if not os.path.exists(img_folder):
        os.makedirs(img_folder)
    # open file to save IMU and odometer data
    imu_path = os.path.join(save_root, 'imu.csv')
    odom_path = os.path.join(save_root, 'odom.csv')
    imu_file = open(imu_path, 'w+')
    imu_file.write('#Timestamp(ns),linearAccX(m/s^2),linearAccY(m/s^2),linearAccZ(m/s^2)'
                   ',AngularVelX(rad/s),AngularVelY(rad/s),AngularVelZ(rad/s)\n')
    odom_file = open(odom_path, 'w+')
    odom_file.write('#Timestamp(ns),linearVelX(m/s),linearVelY(m/s),linearVelZ(m/s)'
                    ',AngularVelX(rad/s),AngularVelY(rad/s),AngularVelZ(rad/s)\n')

    # open ros bag
    bag = rosbag.Bag(bag_file)
    topics = bag.get_type_and_topic_info()[1].keys()
    logging.info('topics in this bag: {}'.format(topics))

    # create cv bridge
    bridge = cv_bridge.CvBridge()

    # read odom topic
    for topic, msg, t in bag.read_messages():
        # print('topic = ', topic, ', t = ', t)

        # time
        stamp = msg.header.stamp
        time_s = msg.header.stamp.secs + 1e-9 * msg.header.stamp.nsecs  # convert to s
        time_ns = 1e9 * msg.header.stamp.secs + msg.header.stamp.nsecs  # convert to ns

        if topic == '/odom':  # Odometer, type: nav_msgs/Odometry
            # type: geometry_msgs/Pose
            position = msg.pose.pose.position
            quaternion = msg.pose.pose.orientation  # [x, y, z, w] with |q| = 1
            # type: geometry_msgs/Twist
            linear = msg.twist.twist.linear  # linear velocity, m/s, [X, Y, Z] = [Forward, Left, Up]
            angular = msg.twist.twist.angular  # angular velocity, rad/s

            logging.debug('odom: t = {:.5f} s, pos = [{:.5g}, {:.5g}, {:.5g}] m'
                          ', quaternion = [{:.5g}, {:.5g}, {:.5g}, {:.5g}], linear = [{:.5g}, {:.5g}, {:.5g}] m/s'
                          ', angular = [{:.5g}, {:.5g}, {:.5g}] rad/s'.format(time_s, position.x, position.y,
                                                                              position.z, quaternion.x, quaternion.y,
                                                                              quaternion.z, quaternion.w, linear.x,
                                                                              linear.y, linear.z, angular.x, angular.y,
                                                                              angular.z))
            # note the |q| = 1, and w != 1
            # print('|q| = {:.5f}'.format(np.sqrt(quaternion.x**2 + quaternion.y**2 + quaternion.z**2 + quaternion.w**2)))

            # write to file
            odom_file.write('{:.0f},{},{},{},{},{},{}\n'.format(time_ns, linear.x, linear.y, linear.z, angular.x,
                                                                angular.y, angular.z))
        elif topic == '/imu/sync':  # IMU, type: sensor_msgs/Imu
            linear_acc = msg.linear_acceleration  #  m/s^2
            angular_vel = msg.angular_velocity  # rad/s
            quaternion = msg.orientation  # [x, y, z, w] with |q| = 1

            logging.debug(
                'imu: t = {:.5f} s, linear acc = [{:.5g}, {:.5g}, {:.5g}] m/s^2'
                ', angular = [{:.5g}, {:.5g}, {:.5g}] m/s, quaternion = [{:.5g}, {:.5g}, {:.5g}, {:.5g}]'.format(
                    time_s, linear_acc.x, linear_acc.y, linear_acc.z, angular_vel.x, angular_vel.y, angular_vel.z,
                    quaternion.x, quaternion.y, quaternion.z, quaternion.w))

            # write to file
            imu_file.write('{:.0f},{},{},{},{},{},{}\n'.format(time_ns, linear_acc.x, linear_acc.y, linear_acc.z,
                                                               angular_vel.x, angular_vel.y, angular_vel.z))
        elif topic == '/sensor_sync_master_node/camera_top/camera_info':  # camera info, type: sensor_msgs/CameraInfo
            # contain image size, camera model, projection and distortion parameters, etc
            width = msg.width
            height = msg.height
            distortion_model = msg.distortion_model
            logging.debug('camera info: t = {:.5f} s, width = {}, height = {}, distortion model = {}'.format(
                time_s, width, height, distortion_model))
        elif topic == '/sensor_sync_master_node/camera_top/image_raw':  # image type: sensor_msgs/Image
            width = msg.width
            height = msg.height
            encoding = msg.encoding
            step = msg.step
            logging.debug('image: t = {:.5f} s, width = {}, height = {}, encoding = {}'.format(
                time_s, width, height, encoding))

            # convert to cv::Mat
            img = bridge.imgmsg_to_cv2(msg)

            # save to file
            img_file = os.path.join(img_folder, '{:.0f}.jpg'.format(time_ns))
            cv.imwrite(img_file, img, (cv.IMWRITE_JPEG_QUALITY, 95))

            # show
            cv.imshow('Image', img)
            cv.waitKey(1)

    # close file
    imu_file.close()
    odom_file.close()


if __name__ == '__main__':
    # config logging
    logging.basicConfig(level=logging.INFO, format="[%(asctime)s %(levelname)s %(filename)s:%(lineno)d] %(message)s")

    # argument parser
    parser = argparse.ArgumentParser(description='Read ROS bag message, and save the sensor data to folder')
    parser.add_argument('--bag', required=True, help='ros bag file')
    parser.add_argument('--save_folder', default='./data/rosbag', help='save folder')
    parser.add_argument('--print_sensor_data', action='store_true', help='save folder')
    args = parser.parse_args()
    print(args)

    if args.print_sensor_data:
        logging.basicConfig(level=logging.DEBUG)

    main(args.bag, args.save_folder)